# coding: utf-8
"""
yolov3 net defines
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from .DarkNet53 import Darknet53
from .block import DBL


class yolov3(nn.Module):
    def __init__(self, class_num):
        super(yolov3, self).__init__()
        self.class_num = class_num
        self.out_ch = 3 * (4 + 1 + class_num)
        # backbone
        self.backbone = Darknet53()

        # head1
        ## DBL是总共6个，其中第5个后，分出来一个给另外一支做上采样了，故先做5个，后再跟一个
        self.head1_1 = nn.Sequential(
            DBL(1024, 512, kernel=1, stride=1, pad=0),
            DBL(512, 1024, kernel=3, stride=1, pad=1),

            DBL(1024, 512, kernel=1, stride=1, pad=0),
            DBL(512, 1024, kernel=3, stride=1, pad=1),

            DBL(1024, 512, kernel=1, stride=1, pad=0)
        )

        self.head1_2 = DBL(512, 1024, kernel=3, stride=1, pad=1)
        self.head1_out = nn.Conv2d(1024, self.out_ch, kernel_size=1, stride=1, padding=0)

        # head2
        ## DBL是总共6个，其中第5个后，分出来一个给另外一支做上采样了，故先做5个，后再跟一个
        self.head2_1_1 = DBL(512, 256, kernel=1, stride=1, pad=0)
        self.head2_2 = nn.Sequential(
            DBL(768, 256, kernel=1, stride=1, pad=0),
            DBL(256, 512, kernel=3, stride=1, pad=1),

            DBL(512, 256, kernel=1, stride=1, pad=0),
            DBL(256, 512, kernel=3, stride=1, pad=1),

            DBL(512, 256, kernel=1, stride=1, pad=0)
        )
        self.head2_3 = DBL(256, 512, kernel=3, stride=1, pad=1)
        self.head2_out = nn.Conv2d(512, self.out_ch, kernel_size=1, stride=1, padding=0)

        # y3
        self.head3_1_1 = DBL(256, 128, kernel=1, stride=1, pad=0)
        ## DBL是总共6个
        self.head3_2 = nn.Sequential(
            DBL(384, 128, kernel=1, stride=1, pad=0),
            DBL(128, 256, kernel=3, stride=1, pad=1),

            DBL(256, 128, kernel=1, stride=1, pad=0),
            DBL(128, 256, kernel=3, stride=1, pad=1),

            DBL(256, 128, kernel=1, stride=1, pad=0),
            DBL(128, 256, kernel=3, stride=1, pad=1),
        )
        self.head3_out = nn.Conv2d(256, self.out_ch, kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        x3, x2, x1 = self.backbone(x)

        # y1
        x1 = self.head1_1(x1)
        y1 = self.head1_2(x1)
        y1 = self.head1_out(y1)

        # y2
        x2_1 = self.head2_1_1(x1)
        x2_1 = F.interpolate(x2_1, scale_factor=2.0)
        x2 = torch.cat([x2, x2_1], dim=1)
        x2 = self.head2_2(x2)
        y2 = self.head2_3(x2)
        y2 = self.head2_out(y2)

        # y3
        x3_1 = self.head3_1_1(x2)
        x3_1 = F.interpolate(x3_1, scale_factor=2.0)
        x3 = torch.cat([x3, x3_1], dim=1)
        x3 = self.head3_2(x3)
        y3 = self.head3_out(x3)

        return y1, y2, y3


if __name__ == "__main__":
    x = torch.randn(1, 3, 416, 416)
    m = yolov3(80)
    y = m(x)
    print("{} {} {}".format(y[0].shape, y[1].shape, y[2].shape))
